'''
Created on Aug 11, 2009

@author: mkiyer
'''

from chipseq.coverage import FixedWidthReadTrack, get_coverage_track
from veggie.io.bowtieparser import BowtieParser
from veggie.genome.chrom import get_chrom_length
import numpy as np

from bx.intervals.intersection import Interval, IntervalTree

import collections
from optparse import OptionParser
import sys
import os
import logging

def read_bed_file(fhd):
    for line in fhd:
        fields = line.strip().split('\t')
        yield (fields[0], int(fields[1]), int(fields[2]))

def extend_intervals(hk36_intervals, dist=10000):
    # extend h3k36 intervals
    h3k36_ext = []
    for chrom, start, end in h3k36_intervals:
        h3k36_ext.append((chrom, max(0, start - dist), min(get_chrom_length(chrom), end + dist)))
    return h3k36_ext  

def pair_nearest_intervals(group1, group2, max_dist):
    '''
    for each interval in group1 find the closest upstream and downstream
    interval in group2
    '''
    # build interval trees for group1 and group2
    trees1 = collections.defaultdict(lambda: IntervalTree())
    trees2 = collections.defaultdict(lambda: IntervalTree())
    for chrom, start, end in group1:
        trees1[chrom].insert_interval(Interval(start, end))
    for chrom, start, end, in group2:         
        trees2[chrom].insert_interval(Interval(start, end))

    pairs = []
    # pair intervals
    for chrom, start1, end1 in group1:

        logging.debug('chrom: %s start: %d end: %d' % (chrom, start1, end1))

        if chrom not in trees2:
            continue
        
        overlapping = trees2[chrom].find(start1, end1)
        if len(overlapping) > 0:
            # if there is more than one interval overlapping, just pick the first one
            # TODO: is this the best option?
            pairs.append([(chrom, start1, end1), (chrom, overlapping[0].start, overlapping[0].end)])
            logging.debug('\toverlapping: %s', overlapping[0])
        else:
            before_match = None
            after_match = None
            before_dist = 0
            after_dist = 0
            
            # not overlapping so try to pair with nearby intervals
            # get the closest interval in group2 which lies before group1
            before = trees2[chrom].before(start1, num_intervals=1, max_dist=max_dist)
            if len(before) > 0:
                # get the closest interval in group1 which lies after the interval in group2
                after = trees1[chrom].after_interval(before[0], num_intervals=1, max_dist=max_dist)
                if len(after) > 0:                    
                    logging.debug('\tbefore group1: %s after group2: %s' % (before[0], after[0]))
                    if (after[0].start, after[0].end) == (start1, end1):
                        before_match = (chrom, before[0].start, before[0].end)
                        before_dist = start1 - before[0].end
            
            # get the closest interval in group1 which lies AFTER group1
            after = trees2[chrom].after(end1, num_intervals=1, max_dist=max_dist)
            if len(after) > 0:
                # get the closest interval in group1 which lies before the interval in group2
                before = trees1[chrom].before_interval(after[0], num_intervals=1, max_dist=max_dist)
                if len(before) > 0:
                    logging.debug('\tafter group1: %s before group2: %s' % (after[0], before[0]))
                    if (before[0].start, before[0].end) == (start1, end1):
                        after_match = (chrom, after[0].start, after[0].end)
                        after_dist = after[0].start - end1

            logging.debug('\tbefore: %s dist: %d' % (before_match, before_dist))
            logging.debug('\tafter: %s dist: %d' % (after_match, after_dist))
            
            if before_match != None and after_match != None:
                # choose the closest of the before and after intervals
                if before_dist < after_dist:
                    pairs.append([(chrom, start1, end1), before_match])
                    logging.debug('\tadded before_match')
                else:
                    pairs.append([(chrom, start1, end1), after_match])
                    logging.debug('\tadded after_match')
            elif before_match != None and after_match == None:                
                pairs.append([(chrom, start1, end1), before_match])
                logging.debug('\tadded before_match')
            elif before_match == None and after_match != None:
                pairs.append([(chrom, start1, end1), after_match])
                logging.debug('\tadded after_match')
            else:
                logging.debug('\tno_match')
    return pairs
    
def intersect(set1, set2):
    result = []
    # build the tree
    intersecters = {}    
    for chrom, start, end in set2:
        if chrom not in intersecters:
            intersecters[chrom] = IntervalTree()
        intersecters[chrom].insert_interval(Interval(start, end, chrom=chrom))
    # intersect
    for chrom, start, end in set1:
        if chrom in intersecters:
            hits = intersecters[chrom].find(start, end)
            if len(hits) == 0:
                continue
            sorted_hits = []
            for hit in hits:
                sorted_hits.append(hit.start)
                sorted_hits.append(hit.end)
            sorted_hits.append(start)
            sorted_hits.append(end)
            sorted_hits.sort()
            result.append((chrom, sorted_hits[0], sorted_hits[-1]))
    return result

def cluster(intervals, cluster_distance, min_size=1):  
    from bx.intervals.cluster import ClusterTree
    # arguments to ClusterTree are:  
    # - Distance in basepairs for two reads to be in the same cluster;  
    #   for instance 20 would group all reads with 20bp of each other  
    # - Number of reads necessary for a group to be considered a cluster;  
    #   2 returns all groups with 2 or more overlapping reads  
    cluster_trees = collections.defaultdict(lambda: ClusterTree(cluster_distance, min_size))    
    id = 0
    for chrom, start, end in intervals:
        cluster_trees[chrom].insert(start, end, id)
        id += 1
    result = []
    for chrom, cluster_tree in cluster_trees.items():
        for start, end, ids in cluster_tree.getregions():
            result.append((chrom, start, end))     
    return result

if __name__ == '__main__':    
    optionparser = OptionParser("usage: %prog [options]")
    optionparser.add_option("--name", dest="name",
                            help="name of experiment")
    optionparser.add_option("--h3k36", dest="h3k36",
                            help="path to h3k36 BED file")
    optionparser.add_option("--h3k4", dest="h3k4",
                            help="path to h3k4 BED file")
    optionparser.add_option("-o", "--output", dest="outfile",
                            default=None,
                            help="output file [default: %default]")    
    (options, args) = optionparser.parse_args()
    if options.outfile == None:
        options.outfile = sys.stdout

    # setup logging
    logging.basicConfig(filename='h3kdetect.log',level=logging.DEBUG)    
    # TODO: check command-line args
    # get the files
    exper_name = options.name
    h3k36_intervals = list(read_bed_file(open(options.h3k36)))
    h3k4_intervals = list(read_bed_file(open(options.h3k4)))
    # intersect h3k36 with h3k4
    pairs = pair_nearest_intervals(h3k36_intervals, h3k4_intervals, max_dist=10000)
    h3k36 = [x[0] for x in pairs]
    # print pairs
    for chrom, start, end in h3k36:
        options.outfile.write("%s\t%d\t%d\n" % (chrom, start, end))            
    # old method: intersect h3k36 with h3k4
    #h3k36_ext = extend_intervals(h3k36_intervals, 10000)
    #pairs = intersect(h3k36_ext, h3k4_intervals)
    #print len(pairs)    
    # cluster to remove overlaps
    #h3k36_and_h3k4 = cluster(h3k36_and_h3k4, 0)

